"""
Author: '夜微凉'
Date: 2025-01-09 10:17:05
LastEditors: '夜微凉'
LastEditTime: 2025-01-09 13:22:00
FilePath: /lrg/utils/quantize.py
Description: 

"""

import torch
import torch.nn as nn
from .module import QuanModuleMapping, IdentityQuan, LsqQuan
from loguru import logger


def quantizer(default_cfg, this_cfg=None):
    target_cfg = dict(default_cfg)
    if this_cfg is not None:
        for k, v in this_cfg.items():
            target_cfg[k] = v

    if target_cfg["bit"] is None:
        q = IdentityQuan
    elif target_cfg["mode"] == "lsq":
        q = LsqQuan
    else:
        raise ValueError("Cannot find quantizer `%s`", target_cfg["mode"])

    target_cfg.pop("mode")
    return q(**target_cfg)


def find_modules_to_quantize(model, quan_scheduler):
    """只量化卷积层和linear层

    :param _type_ model: _description_
    :param _type_ quan_scheduler: _description_
    :return _type_: _description_
    """
    replaced_modules = dict()
    for name, module in model.named_modules():
        if type(module) in QuanModuleMapping.keys():

            if name in quan_scheduler["excepts"]:
                logger.debug(
                    f"excepts Replace module {name}--{type(module)} with quantized module"
                )
                replaced_modules[name] = QuanModuleMapping[type(module)](
                    module,
                    quan_w_fn=quantizer(
                        quan_scheduler["weight"],
                        quan_scheduler["excepts"][name]["weight"],
                    ),
                    quan_a_fn=quantizer(
                        quan_scheduler["act"], quan_scheduler["excepts"][name]["act"]
                    ),
                )
            else:
                logger.debug(
                    f"Replace module {name} -- node:{type(module)} with quantized module {QuanModuleMapping[type(module)]}"
                )
                replaced_modules[name] = QuanModuleMapping[type(module)](
                    module,
                    quan_w_fn=quantizer(quan_scheduler["weight"]),
                    quan_a_fn=quantizer(quan_scheduler["act"]),
                )
        else:
            logger.warning(
                f"This node {name} is not defined and not quantified"
            )
        # elif name in quan_scheduler["excepts"]:
        #     logger.warning(
        #         f"Cannot find module {name}--{type(module)} in the model, skip it"
        #     )

    return replaced_modules

def replace_module_by_names(model, modules_to_replace):
    def helper(child: torch.nn.Module):
        for n, c in child.named_children():
            if type(c) in QuanModuleMapping.keys():
                for full_name, m in model.named_modules():
                    if c is m:
                        child.add_module(n, modules_to_replace.pop(full_name))
                        break
            else:
                helper(c)

    helper(model)
    return model
